import os
import time
from statistics import mean

import torch
import numpy as np
from network import C3D_model, R2Plus1D_model, R3D_model
import cv2

torch.backends.cudnn.benchmark = True


def CenterCrop(frame, size):
    h, w = np.shape(frame)[0:2]
    th, tw = size
    x1 = int(round((w - tw) / 2.))
    y1 = int(round((h - th) / 2.))

    frame = frame[y1:y1 + th, x1:x1 + tw, :]
    return np.array(frame).astype(np.uint8)


def center_crop(frame):
    frame = frame[8:120, 30:142, :]
    return np.array(frame).astype(np.uint8)


# model_path = 'run/run_10/models/C3D-UAV_epoch-9.pth.tar'
# model_path = 'models/R2Plus1D-UAV2_epoch-39.pth.tar'  # 模型路径
# model_path = 'models/R2Plus1D-UAV2_epoch-39.pth.tar'  # 模型路径
# model_path = 'models/R3D-UAV2_epoch-29.pth.tar'  # 模型路径
# model_path = 'models/R3D-UAV2_epoch-29-1.pth.tar'  # 模型路径

model_path = r'D:\pycharmProj\pytorch-video-recognition\run\run_3\models\C3D-UAV4_best.pth.tar'  # 精度最好
# model_path = r'D:\pycharmProj\pytorch-video-recognition\run\run_12\models\C3D-UAV2_epoch-50.pth.tar'  # 最后一个epoch
video_path = 'data/UAV4/AIRPLANE/V_AIRPLANE_001_0.mp4'  # 待检测视频路径
out_video_path = 'run/run_10/manual_test_out/V_AIRPLANE_001_0_testout.mp4'  # 检测结果路径
label_path = r"./dataloaders/uav4_labels.txt"  # 数据集使用的类别标签文件路径

model_name = "C3D"  # 模型类型
num_classes = 4  # 类别数


def main():
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print("Device being used:", device)
    # 读取标签并加载
    with open(label_path, 'r') as f:
        class_names = f.readlines()  # 1 ApplyEyeMakeup
        f.close()
    # 选择要加载的模型
    if model_name == 'R3D':
        model = R3D_model.R3DClassifier(num_classes=num_classes, layer_sizes=(2, 2, 2, 2), pretrained=True)
    elif model_name == 'C3D':
        model = C3D_model.C3D(num_classes=num_classes)
    elif model_name == 'R2Plus1D':
        model = R2Plus1D_model.R2Plus1DClassifier(num_classes=num_classes, layer_sizes=(2, 2, 2, 2), pretrained=True)
    else:
        raise RuntimeError('未知的模型类别' + model_name)
    checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage)

    model.load_state_dict(checkpoint['state_dict'])
    model.to(device)
    model.eval()

    # init model Pytorch加载预训练权重文件时遇到权重不匹配的问题处理
    # model = C3D_model.C3D(num_classes=4) # 4类为UAV4的种类
    # pretrained_dict = torch.load(model_path,
    #                              map_location=lambda storage, loc: storage)
    # model_dict = model.state_dict()
    # # 重新制作预训练的权重，主要是减去参数不匹配的层，楼主这边层名为“fc”
    # pretrained_dict = {k: v for k, v in pretrained_dict.items() if (k in model_dict and 'fc' not in k)}
    # # 更新权重
    # model_dict.update(pretrained_dict)
    # model.load_state_dict(model_dict)
    #
    # model.to(device)
    # model.eval()

    # 推理视频路径
    video = video_path
    cap = cv2.VideoCapture(video)
    retaining = True

    # out为保存的预测结果视频
    width, height, fps, fourcc, out = None, None, None, None, None

    clip = []  # 4维张量：帧，RGB通道
    avg_time = []
    prob_list = []
    probs_list = []
    while retaining:
        retaining, frame = cap.read()
        if not out:  # 获取视频属性
            width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))  # 获取原视频的宽
            height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))  # 获取原视频的搞
            fps = int(cap.get(cv2.CAP_PROP_FPS))  # 帧率
            fourcc = int(cap.get(cv2.CAP_PROP_FOURCC))  # 维持原视频的编码
            # 视频对象的输出路径及文件名
            out = cv2.VideoWriter(out_video_path, fourcc, fps, (width, height))

        if not retaining and frame is None:
            continue
        tmp_ = center_crop(cv2.resize(frame, (171, 128)))
        tmp = tmp_ - np.array([[[90.0, 98.0, 102.0]]])
        clip.append(tmp)  # 添加最新的一帧（添加后为16帧）
        t0 = time.time()
        if len(clip) == 16:  # 如果长度达到16帧（视频前15帧不做预测）
            # 在此更新最小包络框大小、位置
            inputs = np.array(clip).astype(np.float32)
            inputs = np.expand_dims(inputs, axis=0)
            inputs = np.transpose(inputs, (0, 4, 1, 2, 3))  # 将4轴换到123轴前面
            inputs = torch.from_numpy(inputs)
            inputs = torch.autograd.Variable(inputs, requires_grad=False).to(device)
            with torch.no_grad():  # 模型推理
                outputs = model.forward(inputs)
            probs = torch.nn.Softmax(dim=1)(outputs)
            probs_list.append(probs)
            label = torch.max(probs, 1)[1].detach().cpu().numpy()[0]  # 获取概率最大的类别标签
            # 打印类别信息到视频帧上
            cls_name = class_names[label].split(' ')[-1].strip()
            cv2.putText(frame, cls_name, (20, 20),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6,
                        (0, 0, 255), 1)
            # 打印概率信息到视频帧上
            prob = probs[0][label]  # 获取类别概率
            cv2.putText(frame, "prob: %.4f" % prob, (20, 40),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6,
                        (0, 0, 255), 1)
            prob_list.append({str(cls_name): float(prob)})  # 类别概率统计信息
            clip.pop(0)  # 弹出最早一帧（弹出后为15帧）
        t1 = time.time()
        avg_time.append(t1 - t0)
        out.write(frame)  # 写入保存视频
    print("每张图片平均用时：%f" % mean(avg_time))
    print('prob_list', prob_list)
    uav_prob, bird_prob = list(), list()
    for x in prob_list:
        if x.keys().__contains__('BIRD'):
            bird_prob.append(float(x.get('BIRD')))
        elif x.keys().__contains__('AIRPLANE'):
            uav_prob.append(float(x.get('AIRPLANE')))
        elif x.keys().__contains__('DRONE'):
            uav_prob.append(float(x.get('DRONE')))
        elif x.keys().__contains__('HELICOPTER'):
            uav_prob.append(float(x.get('HELICOPTER')))
        else:
            uav_prob.append(float(x.get('UAV')))
    # uav_prob = [float(x.get('UAV')) for x in prob_list if x.keys().__contains__('UAV')]
    # bird_prob = [float(x.get('BIRD')) for x in prob_list if x.keys().__contains__('UAV')]
    print(
        f'uav len:{len(uav_prob)} \tmax prob:{max(uav_prob) if len(uav_prob) > 0 else 0} '
        f'\tavg prob:{mean(uav_prob) if len(uav_prob) > 0 else 0} '
        f'\tmin prob:{min(uav_prob) if len(uav_prob) > 0 else 0}')
    print(
        f'bird len:{len(bird_prob)} \tmax prob:{max(bird_prob) if len(bird_prob) > 0 else 0} '
        f'\tavg prob:{mean(bird_prob) if len(bird_prob) > 0 else 0}'
        f'\tmin prob:{min(bird_prob) if len(bird_prob) > 0 else 0}')
    # print('probs_list', probs_list)
    out.release()
    cap.release()
    cv2.destroyAllWindows()


if __name__ == '__main__':
    main()
